from numpy import extract
import pandas as pd
from collections import defaultdict
import glob
import json
import os
from tqdm import tqdm

fsd50k_csvs = sorted(glob.glob("../../src/FSD50K/FSD50K.ground_truth/*.csv"))[:2]
mapping_file = "../../src/DESED_AS_label_map.json"
desed_filenames = "../../src/DESED_all_filenames.json"

with open(mapping_file, "r") as f:
    label_map = json.load(f)    

# reverse the label map
reverse_label_map = defaultdict(str)
for key, value in label_map.items():
    for v in value:
        reverse_label_map[v] = key

label_map = reverse_label_map
valid_labels = set(label_map.values())
# load all desed filenames
with open(desed_filenames, "r") as f:
    desed_filenames = set(json.load(f))

# load audioset csvs
extracted_df = pd.DataFrame(columns=["fname", "labels", "mids", "desed_labels"])
for csv in fsd50k_csvs:
    print(f"Processing {csv}...")
    df = pd.read_csv(csv, sep=",")
    df["mids"] = df["mids"].apply(lambda x: x[1:-1].replace(" ", ""))
    df["desed_labels"] = df["mids"].apply(lambda x: ",".join([label_map.get(label) if label_map.get(label) is not None else "other" for label in x.split(",")]))    
    df_valid = df[df["desed_labels"].apply(lambda x: any([label in valid_labels for label in x.split(",")]))]
    print(f"Total rows: {len(df)}, Valid rows: {len(df_valid)}")
    extracted_df = pd.concat([extracted_df, df_valid[["fname", "labels", "mids", "desed_labels"]]], ignore_index=True)
    # print(df_valid.head())
# save the extracted dataframe to csv
os.makedirs("../../src/extracted/", exist_ok=True)
extracted_df.to_csv("../../src/extracted/DESED_FSD50K_label_mapping.csv", index=False)

# extract labels contain "Domestic_sounds_and_home_sounds"
dom_df = pd.DataFrame(columns=["fname", "labels", "mids"])
for csv in fsd50k_csvs:
    print(f"Processing {csv}...")
    df = pd.read_csv(csv, sep=",")
    print(df["labels"][0].split(","))
    df_valid = df[df["labels"].apply(lambda x: "Domestic_sounds_and_home_sounds" in x.split(",") or "Domestic_animals_and_pets" in x.split(","))]
    print(f"Total rows: {len(df)}, Valid rows: {len(df_valid)}")
    dom_df = pd.concat([dom_df, df_valid[["fname", "labels", "mids"]]], ignore_index=True)
os.makedirs("../../src/extracted/", exist_ok=True)
dom_df.to_csv("../../src/extracted/DESED_FSD50K_domestic_sounds.csv", index=False)